package device

import (
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"runtime"
	"sync"
	"vpn/guard/config"
	"vpn/guard/network"
	"vpn/guard/utils"
)

type Endpoint struct {
	dstIP net.IP
	srcIP net.IP
}
type Packet struct {
	packet   []byte
	endpoint Endpoint
	buffer   *[]byte //用来回收到对象池

}
type UDPPacket struct {
	packet  []byte
	addr    *net.UDPAddr
	msgType uint32  // 1 握手请求 2 握手响应
	buffer  *[]byte //用来回收到对象池
}
type Peers struct {
	sync.RWMutex
}

type Device struct {
	PrivateKey SecKey
	AllowedIPs *AllowedIPs
	Tun        *DeviceTunnel
	PreUdpPort int // 预设UDP端口，具体生效端口udpConn.LocalAddr()更准
	Done       chan struct{}
	Wg         *sync.WaitGroup
	pool       *Pool
	peers      map[SecKey]*Peer
	udpConn    *net.UDPConn
	indexTable *IndexTable
	queue      struct {
		handshake chan UDPPacket
		send      chan Packet
		recv      chan UDPPacket
	}
	cookieChecker CookieChecker
}

type DeviceTunnel struct {
	Close func() error
	Write func(data []byte) (int, error)
	Read  func(data []byte) (int, error)
	Name  func() string
	Mtu   func() int
}
type NewDeviceParams struct {
	Tun    *DeviceTunnel
	Config *config.Config
}

func NewDevice(params NewDeviceParams) *Device {

	device := &Device{
		PreUdpPort: params.Config.ListenPort,
		Tun:        params.Tun,
		Done:       make(chan struct{}),
		Wg:         &sync.WaitGroup{},
		AllowedIPs: &AllowedIPs{},

		indexTable: NewIndexTable(),
		peers:      make(map[SecKey]*Peer),
		pool:       NewPool(params.Tun.Mtu() + 200), //这里可能会有个坑，udp和tun数据包大小或许会越界。例如tun设置为1000，如果udp数据包大小为1200，那么就会越界。当然wireguard协议的udp数据包大小不到1500，所以这里设置为1600应该足够了。
		queue: struct {
			handshake chan UDPPacket
			send      chan Packet
			recv      chan UDPPacket
		}{
			handshake: make(chan UDPPacket),
			send:      make(chan Packet, 128), //128 1024都是经验值，不合适自行调整
			recv:      make(chan UDPPacket, 1024),
		},
	}
	var Skey SecKey
	Skey.FromString(params.Config.PrivateKey)
	device.PrivateKey = Skey
	device.setCookieChecker(Skey)
	for _, peer := range params.Config.Peers {
		var Pkey SecKey
		Pkey.FromString(peer.PublicKey)
		peer_ := NewPeer()
		peer_.publicKey = Pkey
		peer_.device = device
		peer_.cookieGenerator.Init(Pkey)
		var PreSharedKey SecKey
		PreSharedKey.FromString(peer.PreSharedKey)
		peer_.preSharedKey = PreSharedKey
		if peer.PersistentKeepalive == 0 {
			peer_.persistentKeepaliveInterval.Store(10)
		} else {
			peer_.persistentKeepaliveInterval.Store(peer.PersistentKeepalive)
		}
		//将 10.2.3.2:51820 转换为 UDPAddr
		var err error
		peer_.endpoint.addr, err = net.ResolveUDPAddr("udp", peer.Endpoint)
		if err != nil {
			fmt.Println("Failed to resolve UDP address:", peer.Endpoint)
			continue
		}
		device.peers[Pkey] = peer_
		for _, allowedIP := range peer.AllowedIPs {
			//将 10.2.3.2/24 转换为 IPNet
			_, ipNet, err := net.ParseCIDR(allowedIP)
			if err != nil {
				fmt.Println("Failed to parse CIDR:", allowedIP)
				continue
			}
			device.AllowedIPs.add(ipNet, peer_)
		}
	}

	return device
}

func (device *Device) UpDevice() {
	select {
	case <-device.Done:
		fmt.Println("Device already closed")
		return
	default:
		tag := network.ValidatePort(device.PreUdpPort)
		if !tag {
			fmt.Printf("UDP port already in use:%d\n", device.PreUdpPort)
			return
		}
		var err error
		device.udpConn, err = network.NewUDPConn(device.PreUdpPort)
		if err != nil {
			fmt.Printf("Failed to create UDP conn:%v\n", err)
			return
		}
		fmt.Println("创建udp连接")
		cpus := runtime.NumCPU()
		fmt.Printf("CPUs: %d\n", cpus)
		for i := 0; i < cpus; i++ {
			go device.routineHandshake()
			go device.routineTansportData()
			go device.tansportDataSendHandle()
		}

		go device.handleUDPTraffic()
		go device.handleTUNTraffic()

		//和每个peer建立握手
		device.Wg.Add(1)
		go func() {
			defer device.Wg.Done()
			device.handleHandshake()
		}()
	}
}

func (device *Device) DownDevice() {
	log.Println("Device closing")
	close(device.Done)
	device.udpConn.Close()
	device.Tun.Close()
	device.Wg.Wait()
	log.Println("Device closed ")
}
func (device *Device) setCookieChecker(sk SecKey) {
	device.cookieChecker.Init(sk.publicKey())
}
func (device *Device) lookupPeer(pk SecKey) *Peer {
	peer, ok := device.peers[pk]
	if !ok {
		return nil
	}
	return peer
}

func (device *Device) handleHandshake() {
	for _, peer := range device.peers {
		peer.sendHandshakeInitiation(false)
	}
}
func (device *Device) handleTUNTraffic() {
	device.Wg.Add(1)
	defer device.Wg.Done()
	for {
		select {
		case <-device.Done:
			return
		default:
			buffer := *device.pool.Get()
			n, err := device.Tun.Read(buffer)
			if err != nil {
				fmt.Println("TUN device closed or invalid")
				device.pool.Put(&buffer)
				return
			}
			if n <= 0 {
				device.pool.Put(&buffer)
				continue
			}
			srcIP, dstIP := utils.ParseIPHeader(buffer[:n])
			select {
			case <-device.Done:
				device.pool.Put(&buffer)
				return
			case device.queue.send <- Packet{
				packet: buffer[:n],
				endpoint: Endpoint{
					dstIP: dstIP,
					srcIP: srcIP,
				},
				buffer: &buffer,
			}:
			default:
				device.pool.Put(&buffer)
			}
		}
	}
}
func (device *Device) handleUDPTraffic() {
	device.Wg.Add(1)
	defer device.Wg.Done()
	fmt.Printf("UDP server started on :%d\n", device.udpConn.LocalAddr().(*net.UDPAddr).Port)
	for {
		select {
		case <-device.Done:
			return
		default:
			buffer := *device.pool.Get()
			n, addr, err := device.udpConn.ReadFromUDP(buffer)
			if err != nil {
				fmt.Printf("Read udp error: %v\n", err)
				device.pool.Put(&buffer)
				return
			}
			if n > 0 {
				msgType := binary.LittleEndian.Uint32(buffer[:4])
				data := UDPPacket{
					packet:  buffer[:n],
					addr:    addr,
					msgType: msgType,
					buffer:  &buffer,
				}
				switch msgType {
				case 1, 2:
					select {
					case <-device.Done:
						device.pool.Put(&buffer)
						return
					case device.queue.handshake <- data:
					default:
						device.pool.Put(&buffer)
					}
				case 4:
					select {
					case <-device.Done:
						device.pool.Put(&buffer)
						return
					case device.queue.recv <- data:
					default:
						device.pool.Put(&buffer)
					}
				default:
					device.pool.Put(&buffer)
				}
			}
		}
	}
}
